
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
from dotenv import load_dotenv
import os

load_dotenv()
# 初始化模型
llm = ChatOpenAI(api_key=os.getenv("DASHSCOPE_API_KEY"),
                 base_url=os.getenv("DASHSCOPE_BASE_URL"),
                 model_name="qwen-plus")

# 创建提示模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "您是世界级的技术文档编写者。"),
    ("user", "{input}")
])

# 使用输出解析器
# output_parser = StrOutputParser()
output_parser = JsonOutputParser()


# chain

# 将其添加到上一个链中
chain = prompt | llm | output_parser
# chain = prompt | llm

# 调用它并提出同样的问题。答案是一个字符串，而不是ChatMessage
# 如果你没有让大模型使用json格式输出，会报错
# res = chain.invoke({"input": "LangChain是什么?"})
res = chain.invoke({"input": "LangChain是什么? 问题用question 回答用answer 用JSON格式回复"})

print(res)

